//	Aviv Charikar	Yael Tzirulnikov	305343600	307955609

package modelTraining;

import java.io.*;
import java.util.*;
import java.util.Map.Entry;

public class LidstoneModelTraining {
	private class CountAndProbability {
		public int count;
		public double probability;
		
		public CountAndProbability(int count, double probability) {
			this.count = count;
			this.probability = probability;
		}
	}
	
	private int vocSize;
	private HashMap<String, CountAndProbability> trainingSet;
	private int countTraining;
	private double unseenProb;
	private HashMap<String, Integer> validationSet;
	private int countValidation;
	private double lambda;
	
	public LidstoneModelTraining(int vocSize, String devFilename, int trainSize) throws IOException {
		this(vocSize, devFilename, trainSize, 0);
	}
	
	public LidstoneModelTraining(int vocSize, String devFilename, int trainSize, double lambda) throws IOException {
		this.vocSize = vocSize;
		trainingSet = new HashMap<String, CountAndProbability>();
		countTraining = 0;
		validationSet = new HashMap<String, Integer>();
		countValidation = 0;
		readDevFile(devFilename, trainSize);
		setLambda(lambda);
	}
	
	private void readDevFile(String devFilename, int trainSize) throws IOException {
		BufferedReader reader = new BufferedReader(new FileReader(devFilename));
		String line;
		while ((line = reader.readLine()) != null) {
			// Handle only sentences
			if ((line.length() == 0) || (line.startsWith("<TRAIN")))
				continue;
			
			// Split the words in the sentence and add them to the counts
			String[] lineWords = line.split(" ");
			for (String word : lineWords) {
				if (word.length() > 0) {
					// Add the first words to the training set and the last to the validation set
					if (countTraining < trainSize) {
						CountAndProbability countAndProbability = trainingSet.get(word);
						if (countAndProbability == null)
							trainingSet.put(word, new CountAndProbability(1, 0));
						else
							countAndProbability.count++;
						countTraining++;
					}
					else {
						Integer count = validationSet.get(word);
						if (count == null)
							validationSet.put(word, 1);
						else
							validationSet.put(word, count + 1);
						countValidation++;
					}
				}
			}
		}
		reader.close();
	}
	
	public double getLambda() {
		return lambda;
	}
	
	public void setLambda(double lambda) {
		this.lambda = lambda;
		// Calculate the word probabilities based on the word count in the training and the lambda used for smoothing
		for (Entry<String, CountAndProbability> entry : trainingSet.entrySet())
			entry.getValue().probability = (entry.getValue().count + lambda) / (countTraining + lambda * vocSize);
		// An unseen word has a count of 0
		unseenProb = lambda / (countTraining + lambda * vocSize);
	}
	
	public Set<String> getAllTrainingWords() {
		return trainingSet.keySet();
	}
	
	public int getWordCount(String word) {
		// Return the count of the word or 0 if the word is unseen
		CountAndProbability countAndProbability = trainingSet.get(word);
		if (countAndProbability == null)
			return 0;
		return countAndProbability.count;
	}
	
	public double getWordFrequencyByCount(int count) {
		return (countTraining * ((count + lambda) / (countTraining + lambda * vocSize)));
	}
	
	public double getWordProbability(String word) {
		// Return the probability of the word or the unseen probability if the word is unseen
		CountAndProbability countAndProbability = trainingSet.get(word);
		if (countAndProbability == null)
			return unseenProb;
		return countAndProbability.probability;
	}
	
	public double getUnseenWordProbability() {
		return unseenProb;
	}
	
	public int getCountTraining() {
		return countTraining;
	}
	
	public int getCountValidation() {
		return countValidation;
	}
	
	public double calcPerplexityOnValidationSet() {
		// Calculate the perplexity on the validation set
		double perp = 0;
		for (Entry<String, Integer> entry : validationSet.entrySet())
			perp += entry.getValue() * Math.log(getWordProbability(entry.getKey())) / Math.log(2);
		perp /= countValidation;
		perp = Math.pow(2, -perp);
		return perp;
	}
	
	public double calcPerplexityOnTestSet(BasicModelTraining testSet) {
		// Calculate the perplexity on the test set
		double perp = 0;
		for (String word : testSet.getAllWords()) {
			int wordTestCount = testSet.getWordCount(word);
			perp += wordTestCount * Math.log(getWordProbability(word)) / Math.log(2);
		}
		perp /= testSet.getCountAll();
		perp = Math.pow(2, -perp);
		return perp;
	}
}
